# models/common.py 中添加
import torch
import torch.nn as nn
import torch.nn.functional as F


class WeightedFeatureFusion(nn.Module):
    def __init__(self, in_channels, weight=True):
        super().__init__()
        self.weight = weight
        if weight:
            self.w = nn.Parameter(torch.ones(len(in_channels), dtype=torch.float32), requires_grad=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        if self.weight:
            w = torch.sigmoid(self.w)
            w_sum = torch.sum(w) + 1e-4
            w = w / w_sum

            out = 0
            for i in range(len(x)):
                out += w[i] * x[i]
        else:
            out = torch.stack(x, dim=0).sum(dim=0)

        return self.relu(out)


class BiFPNBlock(nn.Module):
    def __init__(self, channels, levels=3):
        super().__init__()
        self.levels = levels

        # 特征变换卷积
        self.td_convs = nn.ModuleList()
        self.bu_convs = nn.ModuleList()

        # 特征融合层
        self.td_fusions = nn.ModuleList()
        self.bu_fusions = nn.ModuleList()

        # 初始化层
        for i in range(levels - 1):
            self.td_fusions.append(WeightedFeatureFusion([channels, channels], weight=True))
            self.td_convs.append(Conv(channels, channels, 3, 1))

        for i in range(levels - 1):
            self.bu_fusions.append(WeightedFeatureFusion([channels, channels, channels], weight=True))
            self.bu_convs.append(Conv(channels, channels, 3, 1))

    def forward(self, inputs):
        assert len(inputs) == self.levels

        # 保存原始输入，用于残差连接
        original_inputs = inputs.copy()

        # 自顶向下路径
        top_down_features = [inputs[-1]]
        for i in range(self.levels - 2, -1, -1):
            # 上采样
            up_feat = F.interpolate(top_down_features[-1], scale_factor=2, mode='nearest')
            # 特征融合
            td_feat = self.td_fusions[self.levels - 2 - i]([up_feat, inputs[i]])
            # 特征变换
            td_feat = self.td_convs[self.levels - 2 - i](td_feat)
            top_down_features.append(td_feat)

        # 反转得到从低到高的顺序
        top_down_features = top_down_features[::-1]

        # 自底向上路径
        bottom_up_features = [top_down_features[0]]
        for i in range(1, self.levels):
            # 下采样
            down_feat = F.max_pool2d(bottom_up_features[-1], kernel_size=2)
            # 特征融合
            bu_feat = self.bu_fusions[i - 1]([down_feat, top_down_features[i], original_inputs[i]])
            # 特征变换
            bu_feat = self.bu_convs[i - 1](bu_feat)
            bottom_up_features.append(bu_feat)

        return bottom_up_features